package algorithm.ordered_table;

//积压结构
public class SizeBalancedTreeMap {

    public static class SBTNode<K extends Comparable<? super K>, V>{
        private K key;
        private V value;
        private SBTNode<K, V> l;
        private SBTNode<K, V> r;
        private int size;

        public SBTNode(K key, V value){
            this.key = key;
            this.value = value;
            size = 1;
        }
    }

    public static class SBTreeMap<K extends Comparable<? super K>, V>{
        private SBTNode<K, V> root;

        public SBTreeMap(){
            root = null;
        }

        private SBTNode<K, V> rightRotate(SBTNode<K, V> cur){
            SBTNode<K, V> left = cur.l;
            cur.l = left.r;
            left.r = cur;
            left.size = cur.size;
            cur.size = (cur.l == null ? 0 : cur.l.size) + (cur.r == null ? 0 : cur.r.size) + 1;
            return left;
        }

        private SBTNode<K, V> leftRotate(SBTNode<K, V> cur){
            SBTNode<K, V> right = cur.r;
            cur.r = right.l;
            right.l = cur;
            right.size = cur.size;
            cur.size = (cur.l == null ? 0: cur.l.size) + (cur.r == null ? 0 : cur.r.size) + 1;
            return right;
        }

        //SB树保持平衡的策略是当前节点的子节点的规模不能比这个子节点的侄子节点的规模小，
        // 保证了当前节点左右两颗子树的规模控制在两倍以内
        //
        //所以只左旋或者右旋并不能保证其平衡因子的要求，旋转后会有两个/三个节点的子树发生改变，所以要保证这两个/三个节点的平衡性
        private SBTNode<K, V> maintain(SBTNode<K, V> cur){
            if(cur == null)
                return null;
            int leftSize = cur.l == null ? 0 : cur.l.size;
            int rightSize = cur.r == null ? 0 : cur.r.size;
            int leftLeftSize = cur.l == null ? 0 : cur.l.l == null ? 0 : cur.l.l.size;
            int leftRightSize = cur.l == null ? 0 : cur.l.r == null ? 0 : cur.l.r.size;
            int rightLeftSize = cur.r == null ? 0 : cur.r.l == null ? 0 : cur.r.l.size;
            int rightRightSize = cur.r == null ? 0 : cur.r.r == null ? 0 : cur.r.r.size;
            if(rightSize < leftLeftSize){
                cur = rightRotate(cur);
                cur.r = maintain(cur.r);
                cur = maintain(cur);
            }else if(rightSize < leftRightSize){
                cur.l = leftRotate(cur.l);
                cur = rightRotate(cur);
                cur.l = maintain(cur.l);
                cur.r = maintain(cur.r);
                cur = maintain(cur);
            }else if(leftSize < rightRightSize){
                cur = leftRotate(cur);
                cur.l = maintain(cur.l);
                cur = maintain(cur);
            }else if(leftSize < rightLeftSize){
                cur.r = rightRotate(cur.r);
                cur = leftRotate(cur);
                cur.l = maintain(cur.l);
                cur.r = maintain(cur.r);
                cur = maintain(cur);
            }
            return cur;
        }

        private SBTNode<K, V> add(SBTNode<K, V> cur, K key, V value){
            if(cur == null)
                return new SBTNode<>(key, value);
            if(cur.key.compareTo(key) > 0){
                cur.l = add(cur.l, key, value);
            }else {
                cur.r = add(cur.r, key, value);
            }
            cur.size += 1;
            return maintain(cur);
        }

        //在SB树中删除节点时不需要对沿途的每一个节点进行维护平衡性，因为SB树一旦添加了一个节点就会整个树变成平衡的
        private SBTNode<K, V> delete(SBTNode<K, V> cur, K key){
            cur.size--;
            if(cur.key.compareTo(key) > 0)
                cur.l = delete(cur.l, key);
            else if(cur.key.compareTo(key) < 0)
                cur.r = delete(cur.r, key);
            else if(cur.key.compareTo(key) == 0){
                if(cur.l == null && cur.r == null)
                    cur = null;
                else if(cur.l != null && cur.r == null)
                    cur = cur.l;
                else if(cur.l == null && cur.r != null)
                    cur = cur.r;
                else {
                    SBTNode<K, V> tmp = cur.l;
                    while (tmp .r != null)
                        tmp = tmp.r;
                    cur.l = delete(cur.l, tmp.key);
                    cur.key = tmp.key;
                    cur.value = tmp.value;
                }
            }
            return cur;
        }

        private SBTNode<K, V> findLastIndex(K key){
            SBTNode<K, V> pre = root;
            SBTNode<K, V> cur = root;
            while(cur != null){
                pre = cur;
                if(cur.key.compareTo(key) == 0)
                    return cur;
                else if (cur.key.compareTo(key) > 0)
                    cur = cur.l;
                else
                    cur = cur.r;
            }
            return pre;
        }

        private SBTNode<K, V> findLastNoSmallIndex(K key){//找到第一个比他大或者等于他的节点，等于的时候就结束了
            SBTNode<K, V> ans = null;
            SBTNode<K, V> cur = root;
            while(cur != null){
                if(cur.key.compareTo(key) == 0)
                    return cur;
                else if(cur.key.compareTo(key) > 0){
                    ans = cur;
                    cur = cur.l;
                }else {
                    cur = cur.r;
                }
            }
            return ans;
        }

        private SBTNode<K, V> findLastNoBigIndex(K key){//找到最后一个比他小或者等于他的节点
            SBTNode<K, V> ans = null;
            SBTNode<K, V> cur = root;
            while(cur != null){
                if(cur.key.compareTo(key) == 0)
                    return cur;
                else if(cur.key.compareTo(key) < 0){
                    ans = cur;
                    cur = cur.r;
                }else {
                    cur = cur.l;
                }
            }
            return ans;
        }

        //非递归方式实现
        private SBTNode<K, V> getIndex(SBTNode<K, V> cur, int kth){
            int num = (cur.l == null ? 0 : cur.l.size) + 1;
            while(num != kth){
                if(num > kth){
                    cur = cur.l;
                    num -= (cur.r == null ? 0 : cur.r.size) + 1;
                }
                else {
                    cur = cur.r;
                    num += (cur.l == null ? 0 : cur.l.size) + 1;
                }
            }
            return cur;
        }

        //递归方式实现
        private SBTNode<K, V> getIndex2(SBTNode<K, V> cur, int kth) {
            int index = (cur.l == null ? 0 : cur.l.size) + 1;
            if (index == kth) {
                return cur;
            } else if (index > kth) {
                return getIndex2(cur.l, kth);
            } else {
                return getIndex2(cur.r, kth - index);
            }
        }

        public void put(K key, V value){
            if(key != null){
                SBTNode<K, V> lastNode = findLastIndex(key);
                if(lastNode != null && lastNode.key.compareTo(key) == 0)
                    lastNode.value = value;
                else {
                    root = add(root, key, value);
                }
            }
        }

        public void remove(K key){
            if(key != null){
                if(containsKey(key))
                    root = delete(root, key);
            }
        }

        public boolean containsKey(K key){
            if(key != null){
                SBTNode<K, V> lastNode = findLastIndex(key);
                return lastNode != null && lastNode.key.compareTo(key) == 0 ? true : false;
            }
            return false;
        }

        public V get(K key){
            if(key != null){
                SBTNode<K, V> node = findLastIndex(key);
                return node != null && node.key.compareTo(key) == 0 ? node.value : null;
            }
            return null;
        }

        public K getIndexKey(int index){
            if(size() >= index && index >= 1){
               return getIndex2(root, index).key;
            }
            return null;
        }

        public V getIndexValue(int index){
            if(size() >= index && index >= 1){
                return getIndex2(root, index).value;
            }
            return null;
        }

        public int size(){
            return root == null ? 0 : root.size;
        }

    }

    public static void main(String[] args) {
        SBTreeMap<String, Integer> sbt = new SBTreeMap<>();
//        System.out.println(sbt.getIndexKey(1));
        sbt.put("d", 4);
        sbt.put("c", 3);
        sbt.put("a", 1);
        sbt.put("b", 2);
//         sbt.put("e", 5);
        sbt.put("g", 7);
        sbt.put("f", 6);
        sbt.put("h", 8);
        sbt.put("i", 9);
        sbt.put("a", 111);
        System.out.println(sbt.get("a"));
        sbt.put("a", 1);
        System.out.println(sbt.get("a"));
        for (int i = 1; i <= sbt.size(); i++) {
            System.out.println(sbt.getIndexKey(i) + " , " + sbt.getIndexValue(i));
        }
//        System.out.println(sbt.findLastNoBigIndex("e").key);
//        System.out.println(sbt.findLastNoSmallIndex("e").key);
//        sbt.remove("g");
//        System.out.println("===========================");
//        for (int i = 1; i <= sbt.size(); i++) {
//            System.out.println(sbt.getIndexKey(i) + " , " + sbt.getIndexValue(i));
//        }
    }
}
